# -*- coding: utf-8 -*-

from mitmproxy.http import HTTPFlow
from mitmproxy.tcp import TCPFlow
from mitmproxy.udp import UDPFlow
from mitmproxy.dns import DNSFlow
from mitmproxy.flow import Flow
from mitmproxy.tls import ClientHelloData, TlsData
from mitmproxy import ctx
from mitmproxy.connection import Client, Server
from collections.abc import Sequence
from mitmproxy.proxy import layers, layer, server_hooks
from mitmproxy.addonmanager import Loader


class AddonBase:
    def __init__(self):
        pass

    def load(self, loader: Loader):
        '''
        第一次加载插件时调用。此事件接收Loader对象，该对象包含用于添加选项和命令的方法。
        这个方法是插件配置自己的地方。
        :param loader:
        :return:
        '''
        pass

    def running(self):
        '''
        当代理完全启动并运行时调用。在这一点上，你可以期望所有的插件都被加载，所有的选项都被设置。
        :return:
        '''
        pass

    def configure(self, updated: set[str]):
        '''
        当配置改变时调用。更新后的参数是一个类似集合的对象，包含所有更改选项的键。
        此事件在使用更新集中的所有选项启动时调用。
        :param updated:
        :return:
        '''
        pass

    def done(self):
        '''
        当插件关闭时调用，无论是通过从mitmproxy实例中删除，还是当mitmproxy本身关闭。
        在关闭时，此事件在事件循环终止后被调用，从而保证它将是插件看到的最后一个事件。
        请注意，日志处理程序此时已关闭，因此对日志函数的调用将不会产生任何输出。
        :return:
        '''
        pass

    def server_connect(self, data: server_hooks.ServerConnectionHookData):
        '''
        Mitmproxy即将连接服务器。注意，一个连接可以对应多个请求。
        设置data.server.error会终止连接。
        :param data:
        :return:
        '''

    def server_connected(self, data: server_hooks.ServerConnectionHookData):
        '''
        Mitmproxy已连接到服务器。
        :param data:
        :return:
        '''
        pass

    def server_disconnected(self, data: server_hooks.ServerConnectionHookData):
        '''
        服务器连接已关闭(主动或被动断开服务器链接)。
        :param data:
        :return:
        '''
        pass

    def client_connected(self, client: Client):
        '''
        客户端已连接mitmproxy。注意，一个连接可以对应多个HTTP请求。
        设置客户端。错误终止连接。
        :param client:
        :return:
        '''
        pass

    def client_disconnected(self, client: Client):
        '''
        客户端连接已关闭(主动或被动关闭)。
        :param client:
        :return:
        '''
        pass

    def http_connect(self, flow: HTTPFlow):
        '''
        日志含义收到HTTP CONNECT请求。对于大多数实际用途，可以忽略此事件。
        此事件仅发生在常规和上游代理模式下，当客户端指示mitmproxy打开到上游主机的连接时。
        在流上设置非2xx响应将向客户端返回响应并中止连接。
        CONNECT请求是mitmproxy本身的HTTP代理指令，不转发。
        它们不生成通常的HTTP处理程序事件，但是通过新打开的连接的所有请求都会生成。
        :param flow:
        :return:
        '''
        pass

    def http_connect_upstream(self, flow: HTTPFlow):
        '''
        HTTP CONNECT请求即将发送到上游代理。
        此事件可用于为上游代理设置自定义身份验证头。
        CONNECT请求不会生成通常的HTTP处理程序事件，但是通过新打开的连接的所有请求都会生成。
        :param flow:
        :return:
        '''
        pass

    def error(self, flow: HTTPFlow):
        '''
        错误事件
        :param flow:
        :return:
        '''
        pass

    def requestheaders(self, flow: HTTPFlow):
        '''
        成功读取请求头。此时，身体是空的。
        :param flow:
        :return:
        '''
        pass

    def request(self, flow: HTTPFlow):
        '''
        发送请求
        :param flow:
        :return:
        '''

    def responseheaders(self, flow: HTTPFlow):
        '''
        成功读取HTTP响应头。此时，身体是空的。
        :param flow:
        :return:
        '''

    def response(self, flow: HTTPFlow):
        '''
        响应
        :param flow:
        :return:
        '''
        # print("响应", flow.response.get_state())
        pass

    def dns_request(self, flow: DNSFlow):
        '''
        收到DNS请求。处理步骤
        :param flow:
        :return:
        '''
        pass

    def dns_response(self, flow: DNSFlow):
        '''
        已接收或设置DNS响应。
        :param flow:
        :return:
        '''
        pass

    def dns_error(self, flow: DNSFlow):
        '''
        发生DNS错误。
        :param flow:
        :return:
        '''
        pass

    def tcp_start(self, flow: TCPFlow):
        '''
        TCP连接已经启动。
        :param flow:
        :return:
        '''
        pass

    def tcp_message(self, flow: TCPFlow):
        '''
        一个TCP连接收到一条消息。最近的消息将是flow.messages[-1]。该消息可由用户修改。
        :param flow:
        :return:
        '''
        pass

    def tcp_end(self, flow: TCPFlow):
        '''
        TCP连接已断开。
        :param flow:
        :return:
        '''
        pass

    def tcp_error(self, flow: TCPFlow):
        '''
        发生TCP错误。
        每个TCP流将接收tcp_error或tcp_end事件，但不能同时接收两者。
        :param flow:
        :return:
        '''
        pass

    def udp_start(self, flow: UDPFlow):
        '''
        UDP连接已经启动。
        :param flow:
        :return:
        '''
        pass

    def udp_message(self, flow: UDPFlow):
        '''
        UDP连接收到一条消息。最近的消息将是flow.messages[-1]。该消息可由用户修改。
        :param flow:
        :return:
        '''
        pass

    def udp_end(self, flow: UDPFlow):
        '''
        UDP连接已结束。
        :param flow:
        :return:
        '''
        pass

    def udp_error(self, flow: UDPFlow):
        '''
        发生UDP错误。
        每个UDP流将接收udp_error或udp_end事件，但不能同时接收两者。
        :param flow:
        :return:
        '''
        pass

    def websocket_start(self, flow: HTTPFlow):
        '''
        WebSocket连接已经开始。
        :param flow:
        :return:
        '''
        pass

    def tls_clienthello(self, data: ClientHelloData):
        '''
        Mitmproxy收到TLS ClientHello消息。
        这个钩子决定是否需要一个服务器连接来与客户端协商TLS (data.establish_server_tls_first)
        :param data:
        :return:
        '''
        pass

    def tls_start_client(self, data: TlsData):
        '''
        mitmproxy与客户端的TLS协商即将开始。处理步骤
        需要一个插件来初始化data.ssl_conn。(默认情况下，这是由mitmproxy.addons.tlsconfig完成的)
        :param data:
        :return:
        '''
        pass

    def tls_start_server(self, data: TlsData):
        '''
        mitmproxy与服务器之间的TLS协商即将开始。
        需要一个插件来初始化data.ssl_conn。(默认情况下，这是由mitmproxy.addons.tlsconfig完成的)
        :param data:
        :return:
        '''
        pass

    def tls_established_client(self, data: TlsData):
        '''
        与客户端的TLS握手已成功完成。
        :param data:
        :return:
        '''
        pass

    def tls_established_server(self, data: TlsData):
        '''
        与服务器的TLS握手已成功完成。
        :param data:
        :return:
        '''
        pass

    def tls_failed_client(self, data: TlsData):
        '''
        与客户端TLS握手失败。
        :param data:
        :return:
        '''
        pass

    def tls_failed_server(self, data: TlsData):
        '''
        与服务器的TLS握手失败。
        :param data:
        :return:
        '''
        pass

    def websocket_message(self, flow: HTTPFlow):
        '''
        当从客户端或服务器接收到WebSocket消息时调用。最近的消息将是flow.messages[-1]。
        该消息可由用户修改。目前有两种类型的消息，对应于BINARY和TEXT帧类型。
        :param flow:
        :return:
        '''
        pass

    def websocket_end(self, flow: HTTPFlow):
        '''
        WebSocket连接已经结束。你可以检查flow.websocket。Close_code来确定它结束的原因。
        :param flow:
        :return:
        '''
        pass

    def socks5_auth(self, data: layers.modes.Socks5AuthData):
        '''
        Mitmproxy已收到用户名/密码SOCKS5凭据。
        这个钩子通过设置data.valid来决定它们是否有效。
        :param data:
        :return:
        '''
        pass

    def next_layer(self, data: layer.NextLayer):
        '''
        网络层正在切换。你可以通过设置data.layer来改变将要使用的层。
        (默认情况下，这是由mitmproxy.addons.NextLayer完成的)
        :param data:
        :return:
        '''
        pass

    def update(self, flows: Sequence[Flow]):
        '''
        当一个或多个流对象被修改时，通常从不同的插件调用Update。
        :param flows:
        :return:
        '''
        pass

    # 动态取消代理
    def CanelProxy(self, flow: HTTPFlow):
        '''
        :param flow:flow
        :return: bool
        '''
        try:
            self.proxy_address = None
            mode: Sequence[str] = ["regular"]
            mode_option = {'mode': mode}
            ctx.master.options.update(**mode_option)
            if flow.server_conn.via:
                con = Server(flow.server_conn.address)
                con.via = None
                flow.server_conn = con
        except Exception as E:
            print(f"取消Proxy异常：{E}")

    # 动态设置代理
    def SetProxy(self, flow: HTTPFlow, proxy_host, proxy_port, proxy_type="http", username=None, password=None):
        '''
        :param flow: flow
        :param proxy: 设置二级代理格式： http://localhost:7889
        :param username:帐号密码格式：username:password
        :return: bool
        '''

        try:
            if flow.server_conn.via[0] == proxy_type and \
                    flow.server_conn.via[1][0] == proxy_host and \
                    flow.server_conn.via[1][1] == proxy_port:
                return True
        except Exception as E:
            pass
        mode: Sequence[str] = [f'upstream:{proxy_type}://{proxy_host}:{proxy_port}']
        mode_option = {'mode': mode}
        ctx.master.options.update(**mode_option)

        if username and password:
            upstream_auth: Sequence[str] = f"{username}:{password}"
            mode_option = {'upstream_auth': upstream_auth}
            ctx.master.options.update(**mode_option)

        if not flow.server_conn.via:
            con = Server(flow.server_conn.address)
            con.via = (proxy_type, (proxy_host, proxy_port))
            flow.server_conn = con
        return True


if __name__ == '__main__':
    pass
